"""
PCD 转 Voxel 可行域地图。
@author: CeasarSmj
@date: 2025-4-14

算法流程：
1. 读取 PCD 文件,获取点云数据 pc_data
2. 把点云 pc_data 转化为体素图 raw_voxel_data, 并进行滤波,得到 voxel_data
3. 对点云数据进行卷积,第一次为可行区域(本身和上面为空,下面不为空),第二次是区域可行域交叠的情况(两个可行域上下交叠就取下面的);
    得到可行域地图 voxel_map。
4. 裁剪 voxel_map , 裁剪去空白部分,并计算出世界坐标系(以pcd中的为准)相对voxel_map原点坐标系的变换参数[x,y,z,roll,pitch,yaw]。
5. 把 voxel_map 保存为 o3d grid + yaml 格式,yaml中写入分辨率和坐标系参数
6. 用 matplotlib 进行可视化

体素地图中,有障碍物(点云)为1,没有为0。
可行域地图中,可行区域为1,非可行区域为0。
"""

import os
import torch
import numpy as np
import open3d as o3d

import yaml
import matplotlib.pyplot as plt


class VoxelMapGenerator:
    # params
    raw_voxel_resolution = 0.05
    voxel_map_resolution = 0.20

    gaussion_kernel_size = 100
    gaussion_sigma = 1

    # 可行域的计算方式是做卷积,然后查看是否大于阈值
    passable_kernel_radius = 2
    passable_kernel_weight_function = None
    passable_kernel = None
    passable_threshold = 0

    passable_kernel2 = None

    def __init__(
        self,
        raw_voxel_resolution=0.05,
        voxel_map_resolution=0.20,
        gaussion_kernel_size=-1,
        gaussion_sigma=0.1,
        passable_kernel_radius=0.5,
        passable_kernel_threshold=0,
        passable_kernel_weight_function=None,
        enable_step_visualization=False,  # 是否启用逐步可视化
        enable_final_visualization=True,  # 是否启用最终结果可视化
    ):
        # 初始化参数
        self.raw_voxel_resolution = raw_voxel_resolution
        self.voxel_map_resolution = voxel_map_resolution
        self.gaussion_kernel_size = (
            gaussion_kernel_size
            if gaussion_kernel_size > 0
            else int(0.5 / raw_voxel_resolution)
        )
        self.gaussion_sigma = gaussion_sigma
        self.passable_kernel_radius = passable_kernel_radius
        self.passable_kernel_threshold = passable_kernel_threshold
        self.enable_step_visualization = enable_step_visualization  # 保存逐步可视化开关
        self.enable_final_visualization = (
            enable_final_visualization  # 保存最终结果可视化开关
        )
        # cuda 加速
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"使用设备: {self.device}")

        # 设置更合理的kernel参数
        self.kernel_radius_N = max(
            1, int(self.passable_kernel_radius / self.voxel_map_resolution)
        )
        print(f"可行域核半径: {self.kernel_radius_N} 体素")
        # 体素地图的维度分布为 (x, y, z)
        # 初始化可行域卷积核
        N = lambda x: np.exp(-(x**2))  # 高斯分布函数
        sigmoid = lambda x: 1 / (1 + np.exp(-x))

        self.passable_kernel_weight_function = (
            (
                lambda r, h: N(r) * ((sigmoid(-h) * 4 - sigmoid(h) * 2) - N(h) * 5)
            )  # 默认可行域为本身的位置为空,上方尽量为空且下方不为空,这样得分应该比较高
            if passable_kernel_weight_function is None
            else passable_kernel_weight_function
        )

        # 创建passable_kernel
        self.passable_kernel = [
            [
                [
                    self.passable_kernel_weight_function(np.sqrt(x**2 + y**2), z)
                    for z in range(-self.kernel_radius_N, self.kernel_radius_N + 1)
                ]
                for y in range(-self.kernel_radius_N, self.kernel_radius_N + 1)
            ]
            for x in range(-self.kernel_radius_N, self.kernel_radius_N + 1)
        ]
        self.passable_kernel = torch.tensor(
            self.passable_kernel, dtype=torch.float32
        ).to(self.device)
        self.passable_kernel2 = torch.tensor([[[1, -1]]], dtype=torch.float32).to(
            self.device
        )

        # 可视化设置
        plt.rcParams["font.sans-serif"] = ["SimHei"]  # 用来正常显示中文标签
        plt.rcParams["axes.unicode_minus"] = False  # 用来正常显示负号

        print("初始化完成")

    def visualize_3d_data(self, data, title="3D Visualization", max_points=20000):
        """通用的3D数据可视化函数
        Args:
            data: 3D numpy数组或tensor
            title: 图表标题
            max_points: 最大显示点数,超过此数量将进行降采样
        """
        print(f"\n准备可视化: {title}")

        # 确保数据是numpy数组
        if torch.is_tensor(data):
            data = data.detach().cpu().numpy()

        print(f"数据形状: {data.shape}")

        # 计算当前体素数量
        n_voxels = np.sum(data > 0)
        print(f"体素数量: {n_voxels}")

        if n_voxels == 0:
            print("警告：没有可显示的体素")
            return

        # 如果体素数量过大,进行降采样
        if n_voxels > max_points:
            print(f"体素数量过大,进行降采样...")
            # 计算降采样率
            sample_rate = max_points / n_voxels
            # 创建随机mask进行降采样
            mask = np.random.random(data.shape) < sample_rate
            sampled_data = np.zeros_like(data)
            sampled_data[mask] = data[mask]
            data = sampled_data
            print(f"降采样后体素数量: {np.sum(data > 0)}")

        try:
            print("正在创建图形...")
            fig = plt.figure(figsize=(10, 8))
            ax = fig.add_subplot(111, projection="3d")

            print("正在渲染体素...")
            # 获取非零元素的坐标
            x, y, z = np.nonzero(data)

            # 使用scatter3D代替voxels进行更快的渲染
            ax.scatter3D(x, y, z, c="b", marker="s", alpha=0.6)

            ax.set_title(title)
            ax.set_xlabel("X轴")
            ax.set_ylabel("Y轴")
            ax.set_zlabel("Z轴")

            print("显示图形...")
            plt.show()

        except Exception as e:
            print(f"可视化过程出错: {str(e)}")

        print("可视化完成\n")

    def __call__(self, file_path, output_file):
        if os.path.exists(file_path) is False:
            print(f"文件{file_path}不存在")
            return

        # 1. 加载点云数据
        points = self.load_pcd(file_path)
        if len(points) == 0:
            print("错误：点云数据为空")
            return
        print(f"处理文件: {file_path}")
        print(f"点云数据点数: {len(points)}")

        # 2. 生成原始体素地图
        raw_voxel, origin = self.generate_voxel_map(points)
        n_voxels = np.sum(raw_voxel > 0)
        print(f"原始体素地图大小: {raw_voxel.shape}")
        print(f"原始体素数量: {n_voxels}")
        if n_voxels == 0:
            print("错误：生成的原始体素地图为空")
            return
        if self.enable_step_visualization:
            self.visualize_3d_data(raw_voxel, "原始体素地图")

        # 3. 滤波处理
        filtered = self.filter_voxel_map(raw_voxel)
        n_voxels = np.sum(filtered > 0)
        print(f"滤波后体素数量: {n_voxels}")
        if n_voxels == 0:
            print("错误：滤波后体素地图为空")
            return
        if self.enable_step_visualization:
            self.visualize_3d_data(filtered, "滤波后的体素地图")

        # 4. 转换为可行域地图
        passable = self.voxel_map_to_passable_map(filtered)
        n_voxels = np.sum(passable > 0)
        print(f"可行域体素数量: {n_voxels}")
        if n_voxels == 0:
            print("错误：生成的可行域地图为空")
            return
        if self.enable_step_visualization:
            self.visualize_3d_data(passable, "可行域地图")

        # 5. 裁剪和变换
        cropped, transform = self.crop_and_calculate_transform(
            passable, origin, self.voxel_map_resolution
        )
        n_voxels = np.sum(cropped > 0)
        print(f"裁剪后体素数量: {n_voxels}")
        if n_voxels == 0:
            print("错误：裁剪后的地图为空")
            return

        # 最终结果可视化
        if self.enable_final_visualization:
            self.visualize_3d_data(cropped, "最终可行域地图")

        # 6. 保存结果
        self.save_voxel_map(cropped, transform, output_file)
        print(f"体素地图保存完成，保存至{output_file}")
        print(f"最终地图大小: {cropped.shape}")
        print(f"最终体素数量: {n_voxels}")

    def load_pcd(self, file_path):
        pcd = o3d.io.read_point_cloud(file_path)
        return np.asarray(pcd.points)

    def generate_voxel_map(self, pc_data):
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(pc_data)
        voxel_grid = o3d.geometry.VoxelGrid.create_from_point_cloud(
            pcd, self.raw_voxel_resolution
        )
        voxels = voxel_grid.get_voxels()
        if not voxels:
            return np.zeros((0, 0, 0), np.zeros(3))
        indices = np.array([voxel.grid_index for voxel in voxels])
        min_indices = indices.min(axis=0)
        max_indices = indices.max(axis=0)
        shape = max_indices - min_indices + 1
        voxel_array = np.zeros(shape, dtype=np.uint8)
        adjusted_indices = indices - min_indices
        voxel_array[
            adjusted_indices[:, 0], adjusted_indices[:, 1], adjusted_indices[:, 2]
        ] = 1
        pcd_min = pc_data.min(axis=0)
        origin = pcd_min + min_indices * self.raw_voxel_resolution
        return voxel_array, origin

    def create_gaussian_kernel_3d(self):
        """创建3D高斯核，形状为 (size, size, size//2)"""
        kernel_size = self.gaussion_kernel_size
        sigma = self.gaussion_sigma

        # 创建更小的高斯核
        x = torch.linspace(-2, 2, kernel_size, dtype=torch.float32)
        y = torch.linspace(-2, 2, kernel_size, dtype=torch.float32)
        z = torch.linspace(-2, 2, kernel_size // 2, dtype=torch.float32)

        x, y, z = torch.meshgrid(x, y, z, indexing="ij")

        # 计算3D高斯核
        kernel = torch.exp(-(x.pow(2) + y.pow(2) + z.pow(2)) / (2 * sigma**2))

        # 归一化
        kernel = kernel / kernel.sum()

        return kernel.to(dtype=torch.float32, device=self.device)

    def filter_voxel_map(self, voxel_map):
        """对体素地图进行滤波处理"""
        kernel = self.create_gaussian_kernel_3d()
        # 确保输入tensor使用float32类型
        tensor = torch.from_numpy(voxel_map).to(dtype=torch.float32, device=self.device)
        tensor = tensor.unsqueeze(0).unsqueeze(0)

        # 打印调试信息
        print(f"\n滤波前体素数量: {torch.sum(tensor > 0).item()}")
        print(f"滤波前tensor形状: {tensor.shape}")

        # 执行高斯滤波
        padding = self.gaussion_kernel_size // 2
        filtered = torch.nn.functional.conv3d(
            tensor, kernel.unsqueeze(0).unsqueeze(0), padding=padding
        )

        print(f"高斯滤波后最大值: {torch.max(filtered).item()}")
        print(f"高斯滤波后最小值: {torch.min(filtered).item()}")

        # 使用较低的动态阈值,取最大值的5%作为阈值
        threshold = torch.max(filtered) * 0.05
        filtered = (filtered > threshold).float()

        print(f"阈值化后体素数量: {torch.sum(filtered > 0).item()}")

        # 降采样,使用更低的密度阈值
        # 计算降采样倍数
        scale_factor = self.voxel_map_resolution / self.raw_voxel_resolution
        kernel_size = int(scale_factor)

        # 使用平均池化进行降采样
        filtered = torch.nn.functional.avg_pool3d(
            filtered, kernel_size=kernel_size, stride=kernel_size
        )
        filtered = (filtered > 0.01).float()  # 降低阈值到1%

        print(f"降采样后体素数量: {torch.sum(filtered > 0).item()}")
        print(f"降采样后tensor形状: {filtered.shape}")
        print(f"降采样倍数: {kernel_size}x\n")

        return filtered.squeeze().cpu().numpy().astype(np.uint8)

    def voxel_tensor_conv3d(self, tensor, kernel):
        """执行3D卷积,确保输入tensor和kernel的数据类型一致"""
        # 确保tensor使用float32类型
        tensor = tensor.to(dtype=torch.float32)
        # 确保kernel使用float32类型
        kernel = kernel.to(dtype=torch.float32)

        tensor = tensor.unsqueeze(0).unsqueeze(0)
        kernel = kernel.unsqueeze(0).unsqueeze(0)

        # 计算padding以保持输出尺寸与输入相同
        padding = tuple((k - 1) // 2 for k in kernel.shape[2:])

        return torch.nn.functional.conv3d(tensor, kernel, padding=padding).squeeze()

    def voxel_map_to_passable_map(self, voxel_map):
        """将体素地图转换为可行域地图，确保卷积操作后的尺寸一致"""
        # 将输入转换为tensor并移至设备
        tensor = torch.from_numpy(voxel_map).float().to(self.device)

        # 第一次卷积：检测可行区域
        conv1 = self.voxel_tensor_conv3d(tensor, self.passable_kernel)
        passable = (conv1 > self.passable_threshold).float()

        # 可视化第一次卷积结果
        if self.enable_step_visualization:
            self.visualize_3d_data(passable.cpu().numpy(), "第一次卷积后的可行域")

        # 第二次卷积：检测上方可行区域
        conv2 = torch.zeros_like(passable)
        if passable.size(2) > 1:  # 确保有足够的高度进行上方检测
            # 对每个z层进行处理，除了最后一层
            conv2[:, :, :-1] = passable[:, :, 1:]

        # 可视化第二次卷积结果
        if self.enable_step_visualization:
            self.visualize_3d_data(conv2.cpu().numpy(), "上方空间检测结果")

        # 应用mask：只在两个区域都有值的地方置1
        mask = (passable > 0) & (conv2 > 0)
        result = torch.zeros_like(passable)
        result[mask] = 1

        return result.cpu().numpy().astype(np.uint8)

    def crop_and_calculate_transform(self, voxel_map, origin, resolution):
        non_zero = np.argwhere(voxel_map)
        if non_zero.size == 0:
            return voxel_map, (0, 0, 0, 0, 0, 0)
        min_coords = non_zero.min(axis=0)
        max_coords = non_zero.max(axis=0)
        cropped = voxel_map[
            min_coords[0] : max_coords[0] + 1,
            min_coords[1] : max_coords[1] + 1,
            min_coords[2] : max_coords[2] + 1,
        ]
        translation = origin + min_coords * resolution
        return cropped, (translation[0], translation[1], translation[2], 0, 0, 0)

    def save_voxel_map(self, voxel_map, transform, output_path):
        # 创建voxel grid
        voxel_grid = o3d.geometry.VoxelGrid()
        voxel_points = np.argwhere(voxel_map)
        voxel_points = voxel_points * self.voxel_map_resolution + np.array(
            transform[:3]
        )
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(voxel_points)
        voxel_grid = o3d.geometry.VoxelGrid.create_from_point_cloud(
            pcd, voxel_size=self.voxel_map_resolution
        )

        # 保存voxel grid
        o3d.io.write_voxel_grid(output_path, voxel_grid)

        # 保存YAML
        yaml_data = {
            "resolution": float(self.voxel_map_resolution),
            "x": float(transform[0]),
            "y": float(transform[1]),
            "z": float(transform[2]),
            "roll": float(transform[3]),
            "pitch": float(transform[4]),
            "yaw": float(transform[5]),
        }
        yaml_path = output_path.replace(".ply", ".yaml")
        with open(yaml_path, "w") as f:
            yaml.dump(yaml_data, f)


if __name__ == "__main__":
    path = "pcds/output.pcd"
    generator = VoxelMapGenerator(
        enable_step_visualization=False, enable_final_visualization=False
    )
    generator(path, "voxel_maps")
